from spk.core.workspace import *
import streamlit as st

from .libs.zip_tree import get_zip_structure

class StepZipTree(SpkWorker):

    '''查看 zip 文件内容

    这是一个简单的工作流
    1. 提供 zip 文件
    2. 以类似 tree 的方式预览结构
    '''

    First = '第一步，提供 zip 文件'
    Second = '第二步，查看文件内容'

    __steps__ = [
        First, Second
    ]

    __keywords__ = [
        'spk-zip-file',
        'spk-zip-tree'
    ]

    def __do_work__(self):

        # 第一步，提供 zip 文件
        if current() == self.First:
            # 如果 zip-file 为空则显示上传
            if get(self.__keywords__[0]) is None:
                if upload_file := st.file_uploader('选择zip文件', accept_multiple_files=False, type=['zip']):
                    st.write(f'准备上传文件: {upload_file.name}')
                    set(self.__keywords__[0], get('workdir') + '/' + upload_file.name)
                    with st.spinner('正在保存文件...'):
                        with open(get(self.__keywords__[0]), 'wb') as file:
                            file.write(upload_file.getvalue())
                            st.rerun()

            # 否则直接显示文件内容
            else:
                next_step(True)

        # 第二步，查看文件内容
        if current() == self.Second:
            
            # 如果 zip-tree 为空则加载
            if get(self.__keywords__[1]) is None:
                tree = get_zip_structure(get(self.__keywords__[0]))
                set(self.__keywords__[1], tree)
                st.rerun()

            # 然后显示 zip-tree 内容
            else:
                st.code(get(self.__keywords__[1]))
    
            if st.button('上一步'):
                # 如果要上一步则需要处理一些状态：
                # 1. 将 zip-file 置空
                # 2. 将 zip-tree 置空
                set(self.__keywords__[0], None)
                set(self.__keywords__[1], None)
                prev_step()